import cv2
import argparse
from rknn_executor import RKNN_model_container
from det import Detection
import numpy as np
import os
from PIL import Image, ImageDraw, ImageFont

# 需要结合目标检测，将车牌检测裁剪后再进行车牌识别，也需要增加矫正，对比度等等操作

class Recognition():
    
    def __init__(self, model_path, nup_core):

        self.CHARS = ['京', '沪', '津', '渝', '冀', '晋', '蒙', '辽', '吉', '黑',
            '苏', '浙', '皖', '闽', '赣', '鲁', '豫', '鄂', '湘', '粤',
            '桂', '琼', '川', '贵', '云', '藏', '陕', '甘', '青', '宁',
            '新',
            '0', '1', '2', '3', '4', '5', '6', '7', '8', '9',
            'A', 'B', 'C', 'D', 'E', 'F', 'G', 'H', 'J', 'K',
            'L', 'M', 'N', 'P', 'Q', 'R', 'S', 'T', 'U', 'V',
            'W', 'X', 'Y', 'Z', 'I', 'O', '-'
            ]
        self.model = RKNN_model_container(model_path,nup_core)
        
    def predict(self, img):
        img = cv2.resize(img, (94, 24))
        img = np.expand_dims(img, 0)
        outputs = self.model.run([img])
        labels, pred_labels = self.decode(outputs[0])
        print('车牌识别结果: '+ labels[0])
        return labels[0]

    def decode(self,preds):
        pred_labels = list()
        labels = list()
        for i in range(preds.shape[0]):
            pred = preds[i, :, :]
            pred_label = list()
            for j in range(pred.shape[1]):
                pred_label.append(np.argmax(pred[:, j], axis=0))
            no_repeat_blank_label = list()
            pre_c = pred_label[0]
            for c in pred_label: 
                if (pre_c == c) or (c == len(self.CHARS) - 1):
                    if c == len(self.CHARS) - 1:
                        pre_c = c
                    continue
                no_repeat_blank_label.append(c)
                pre_c = c
            pred_labels.append(no_repeat_blank_label)
        
        for i, label in enumerate(pred_labels):
            lb = ""
            for i in label:
                lb += self.CHARS[i]
            labels.append(lb)
        return labels, pred_labels

if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='Process some integers.')
    parser.add_argument('--rec_model_path', default=r'./model/rec.rknn', type=str,  help='rec_model_path model path')
    parser.add_argument('--det_model_path', default=r'./model/det.rknn', type=str,  help='rknn model det_model_path')
    parser.add_argument('--img_path', type=str, default=r'./image/222.png', help='image folder path')
    args = parser.parse_args()

    image = cv2.imread(args.img_path)

    if image is None:
        print("image not found")
        exit(0)
    
    detection = Detection(args.det_model_path, 1)
    
    recognition = Recognition(args.rec_model_path,0)
    
    boxes,classes,scores = detection.predict(image)

    if boxes is not None:
        
        pil_img = Image.fromarray(cv2.cvtColor(image, cv2.COLOR_BGR2RGB))
        draw = ImageDraw.Draw(pil_img)
        font = ImageFont.truetype('lib/Arial.Unicode.ttf', 30)
        
        for box, score, cl in zip(boxes, scores, classes):
            # 裁剪车牌 y1:y2, x1:x2
            img = image[int(box[1]):int(box[3]), int(box[0]):int(box[2])]

            # 车牌增强(对比度，锐化, 亮度，高斯等等)

            # 倾斜矫正

            # 透视变换
            
            # 检测车牌
            plate = recognition.predict(img)
            top, left, right, bottom = [int(_b) for _b in box]
            print("%s @ (%d %d %d %d) %.3f" % (detection.CLASSES[cl], top, left, right, bottom, score))
            draw.rectangle([(top, left), (right, bottom)],fill=None , outline=(255, 0, 0),width=2)
            draw.text((top, bottom), plate, font = font, fill=(255, 0, 0))

        opencv_img = cv2.cvtColor(np.array(pil_img), cv2.COLOR_RGB2BGR)
            
    if not os.path.exists('./result'):
        os.mkdir('./result')
    result_path = os.path.join('./result', 'result.png')
    cv2.imwrite(result_path, opencv_img)
    print('推理结果保存在 {}，刷新目录查看'.format(result_path))
